#ifndef TRIL_N_D_H
#define TRIL_N_D_H

#include "kernel_operator.h"

namespace Tril {
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t BYTE_BLOCK = 32;

constexpr float ZERO_FP32 = 0.0;
constexpr float ONE_FP32 = 1.0;

template <typename T>
class TrilND {
public:
    __aicore__ inline TrilND() = default;
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, const TrilTilingData* __restrict tilingData);
    template <bool rowSmall = true>
    __aicore__ inline void ProcessMultiGroup();
    template <bool rowSmall = true>
    __aicore__ inline void ProcessOneGroup(int64_t gmOffset = 0);

private:
    template <typename T1, typename T2>
    __aicore__ inline T1 CeilDiv(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : (a + bTemp - 1) / bTemp;
    };

    template <typename T1, typename T2>
    __aicore__ inline T1 CeilAlignA2B(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : CeilDiv(a, bTemp) * bTemp;
    };

    __aicore__ inline void CopyIn(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void Compute(int64_t dataCount, int64_t dupCount);
    __aicore__ inline void CopyOut(int64_t gmOffset, int64_t dataCount);

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> xQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;
    GlobalTensor<T> xGM, yGM;
    int64_t blockIdx = 0;
    uint64_t perBlockCount = 0;
    uint64_t colCountAlign = 0;
    uint64_t innerTailCountAlign = 0;

    // tiling params
    uint64_t ubMaxProcCount = 0;
    uint64_t totalDataCount = 0;
    uint64_t groupCount = 0;
    uint64_t rowCount = 0;
    uint64_t colCount = 0;
    uint64_t innerLoopCount = 0;
    uint64_t innerTailCount = 0;
    int64_t diagonal = 0;
};

template <typename T>
__aicore__ inline void TrilND<T>::Init(GM_ADDR x, GM_ADDR y, const TrilTilingData* __restrict tilingData) {
    blockIdx = GetBlockIdx();
    perBlockCount = BYTE_BLOCK / sizeof(T);
    ubMaxProcCount = tilingData->ubMaxProcCount;
    totalDataCount = tilingData->totalDataCount;
    groupCount = tilingData->groupCount;
    rowCount = tilingData->rowCount;
    colCount = tilingData->colCount;
    innerLoopCount = tilingData->innerLoopCount;
    innerTailCount = tilingData->innerTailCount;
    diagonal = tilingData->diagonal;

    colCountAlign = CeilAlignA2B(colCount, perBlockCount);
    innerTailCountAlign = CeilAlignA2B(innerTailCount, perBlockCount);

    xGM.SetGlobalBuffer((__gm__ T*)x, totalDataCount);
    yGM.SetGlobalBuffer((__gm__ T*)y, totalDataCount);

    uint64_t singleBufferSize = ubMaxProcCount * sizeof(T);
    pipe.InitBuffer(xQueue, BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(yQueue, BUFFER_NUM, singleBufferSize);
}

template <typename T>
template <bool rowSmall>
__aicore__ inline void TrilND<T>::ProcessMultiGroup() {
    int64_t gmOffset = 0;
    for (int64_t i = 0; i < groupCount; i++) {
        ProcessOneGroup<rowSmall>(gmOffset);
        gmOffset += rowCount * colCount;
    }
}

template <typename T>
template <bool rowSmall>
__aicore__ inline void TrilND<T>::ProcessOneGroup(int64_t gmOffset /*= 0*/) {
    int64_t dupCount = diagonal + 1;
    for (int64_t j = 0; j < rowCount; j++) {
        if constexpr (rowSmall) {
            CopyIn(gmOffset, colCountAlign);
            Compute(colCountAlign, dupCount);
            CopyOut(gmOffset, colCountAlign);
            gmOffset += colCount;
        } else {
            int64_t tmpDupCount = dupCount;
            for (int64_t k = 0; k < innerLoopCount; k++) {
                CopyIn(gmOffset, ubMaxProcCount);
                Compute(ubMaxProcCount, tmpDupCount > ubMaxProcCount ? ubMaxProcCount : tmpDupCount);
                CopyOut(gmOffset, ubMaxProcCount);
                tmpDupCount -= ubMaxProcCount;
                gmOffset += ubMaxProcCount;
            }
            if (innerTailCount) {
                CopyIn(gmOffset, innerTailCountAlign);
                Compute(innerTailCountAlign, tmpDupCount);
                CopyOut(gmOffset, innerTailCountAlign);
                gmOffset += innerTailCount;
            }
        }

        dupCount++;
        dupCount = dupCount > colCount ? colCount : dupCount;
    }
}

template <typename T>
__aicore__ inline void TrilND<T>::CopyIn(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> xInLT = xQueue.AllocTensor<T>();
    DataCopy(xInLT, xGM[gmOffset], dataCount);
    xQueue.EnQue(xInLT);
}

template <typename T>
__aicore__ inline void TrilND<T>::Compute(int64_t dataCount, int64_t dupCount) {
    LocalTensor<T> xLT = xQueue.DeQue<T>();
    LocalTensor<T> yLT = yQueue.AllocTensor<T>();
    Duplicate(yLT, T(ZERO_FP32), dataCount);
    if (dupCount > 0) {
        Duplicate(yLT, T(ONE_FP32), dupCount);
    }
    Mul(yLT, yLT, xLT, dataCount);

    yQueue.EnQue(yLT);
    xQueue.FreeTensor(xLT);
}

template <typename T>
__aicore__ inline void TrilND<T>::CopyOut(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> yOutLT = yQueue.DeQue<T>();
    DataCopy(yGM[gmOffset], yOutLT, dataCount);
    yQueue.FreeTensor(yOutLT);
}
}  // namespace Tril

#endif  // TRIL_N_D_H